{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Content under Creative Commons Attribution license CC-BY 4.0, code under BSD 3-Clause License © 2018  by D. Koehn, heterogeneous models are from [this Jupyter notebook](https://nbviewer.jupyter.org/github/krischer/seismo_live/blob/master/notebooks/Computational%20Seismology/The Finite-Difference Method/fd_ac2d_heterogeneous.ipynb) by Heiner Igel ([@heinerigel](https://github.com/heinerigel)), Florian Wölfl and Lion Krischer ([@krischer](https://github.com/krischer)) which is a supplemenatry material to the book [Computational Seismology: A Practical Introduction](http://www.computational-seismology.org/), notebook style sheet by L.A. Barba, N.C. Clementi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to load the notebook's style sheet, then ignore it\n",
    "from IPython.core.display import HTML\n",
    "css_file = '../../style/custom.css'\n",
    "HTML(open(css_file, \"r\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D acoustic FD modelling for heterogeneous media\n",
    "\n",
    "So far, we only compared 2D acoustic FD modelling results for homogeneous acoustic media with analytical solutions. Next, we want to model some more interesting, heterogeneous problems. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# Import Libraries \n",
    "# ----------------\n",
    "import numpy as np\n",
    "from numba import jit\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from pylab import rcParams\n",
    "\n",
    "# Ignore Warning Messages\n",
    "# -----------------------\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Definition of initial modelling parameters\n",
    "# ------------------------------------------\n",
    "xmax = 2000.0 # maximum spatial extension of the 1D model in x-direction (m)\n",
    "zmax = xmax   # maximum spatial extension of the 1D model in z-direction (m)\n",
    "dx   = 10.0   # grid point distance in x-direction (m)\n",
    "dz   = dx     # grid point distance in z-direction (m)\n",
    "\n",
    "tmax = 0.75    # maximum recording time of the seismogram (s)\n",
    "dt   = 0.0010 # time step\n",
    "\n",
    "vp0  = 3000.  # P-wave speed in medium (m/s)\n",
    "\n",
    "# acquisition geometry\n",
    "xsrc = 1000.0 # x-source position (m)\n",
    "zsrc = xsrc   # z-source position (m)\n",
    "\n",
    "f0   = 100.0 # dominant frequency of the source (Hz)\n",
    "t0   = 0.1   # source time shift (s)\n",
    "\n",
    "isnap = 2  # snapshot interval (timesteps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True) # use JIT for C-performance\n",
    "def update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz):\n",
    "    \n",
    "    for i in range(1, nx - 1):\n",
    "        for j in range(1, nz - 1):\n",
    "                \n",
    "            d2px[i,j] = (p[i + 1,j] - 2 * p[i,j] + p[i - 1,j]) / dx**2                \n",
    "            d2pz[i,j] = (p[i,j + 1] - 2 * p[i,j] + p[i,j - 1]) / dz**2\n",
    "        \n",
    "    return d2px, d2pz   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define simple absorbing boundary frame based on wavefield damping \n",
    "# according to Cerjan et al., 1985, Geophysics, 50, 705-708\n",
    "def absorb(nx,nz):\n",
    "\n",
    "    FW = 60     # thickness of absorbing frame (gridpoints)    \n",
    "    a = 0.0053\n",
    "    \n",
    "    coeff = np.zeros(FW)\n",
    "    \n",
    "    # define coefficients in absorbing frame\n",
    "    for i in range(FW):    \n",
    "        coeff[i] = np.exp(-(a**2 * (FW-i)**2))\n",
    "\n",
    "    # initialize array of absorbing coefficients\n",
    "    absorb_coeff = np.ones((nx,nz))\n",
    "\n",
    "    # compute coefficients for left grid boundaries (x-direction)\n",
    "    zb=0 \n",
    "    for i in range(FW):\n",
    "        ze = nz - i - 1\n",
    "        for j in range(zb,ze):\n",
    "            absorb_coeff[i,j] = coeff[i]\n",
    "\n",
    "    # compute coefficients for right grid boundaries (x-direction)        \n",
    "    zb=0\n",
    "    for i in range(FW):\n",
    "        ii = nx - i - 1\n",
    "        ze = nz - i - 1\n",
    "        for j in range(zb,ze):\n",
    "            absorb_coeff[ii,j] = coeff[i]\n",
    "\n",
    "    # compute coefficients for bottom grid boundaries (z-direction)        \n",
    "    xb=0 \n",
    "    for j in range(FW):\n",
    "        jj = nz - j - 1\n",
    "        xb = j\n",
    "        xe = nx - j\n",
    "        for i in range(xb,xe):\n",
    "            absorb_coeff[i,jj] = coeff[j]\n",
    "\n",
    "    return absorb_coeff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FD_2D_acoustic code with JIT optimization\n",
    "# -----------------------------------------\n",
    "def FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc):        \n",
    "    \n",
    "    # define model discretization\n",
    "    # ---------------------------\n",
    "\n",
    "    nx = (int)(xmax/dx) # number of grid points in x-direction\n",
    "    print('nx = ',nx)\n",
    "\n",
    "    nz = (int)(zmax/dz) # number of grid points in x-direction\n",
    "    print('nz = ',nz)\n",
    "\n",
    "    nt = (int)(tmax/dt) # maximum number of time steps            \n",
    "    print('nt = ',nt)\n",
    "\n",
    "    isrc = (int)(xsrc/dx)  # source location in grid in x-direction\n",
    "    jsrc = (int)(zsrc/dz)  # source location in grid in x-direction\n",
    "\n",
    "    # Source time function (Gaussian)\n",
    "    # -------------------------------\n",
    "    src  = np.zeros(nt + 1)\n",
    "    time = np.linspace(0 * dt, nt * dt, nt)\n",
    "\n",
    "    # 1st derivative of Gaussian\n",
    "    src  = -2. * (time - t0) * (f0 ** 2) * (np.exp(- (f0 ** 2) * (time - t0) ** 2)) \n",
    "    \n",
    "    # define clip value: 0.1 * absolute maximum value of source wavelet\n",
    "    clip = 0.1 * max([np.abs(src.min()), np.abs(src.max())]) / (dx*dz) * dt**2\n",
    "    \n",
    "    # Define absorbing boundary frame\n",
    "    # -------------------------------    \n",
    "    absorb_coeff = absorb(nx,nz)\n",
    "    \n",
    "    # Define model\n",
    "    # ------------    \n",
    "    vp  = np.zeros((nx,nz))\n",
    "    vp  = model(nx,nz,vp,dx,dz)    \n",
    "    vp2 = vp**2\n",
    "    \n",
    "    # Initialize empty pressure arrays\n",
    "    # --------------------------------\n",
    "    p    = np.zeros((nx,nz)) # p at time n (now)\n",
    "    pold = np.zeros((nx,nz)) # p at time n-1 (past)\n",
    "    pnew = np.zeros((nx,nz)) # p at time n+1 (present)\n",
    "    d2px = np.zeros((nx,nz)) # 2nd spatial x-derivative of p\n",
    "    d2pz = np.zeros((nx,nz)) # 2nd spatial z-derivative of p    \n",
    "    \n",
    "    # Initalize animation of pressure wavefield \n",
    "    # -----------------------------------------    \n",
    "    fig = plt.figure(figsize=(7,3.5))  # define figure size\n",
    "    plt.tight_layout()\n",
    "    extent = [0.0,xmax,zmax,0.0]     # define model extension\n",
    "    \n",
    "    # Plot pressure wavefield movie\n",
    "    ax1 = plt.subplot(121)\n",
    "    image = plt.imshow(p.T, animated=True, cmap=\"RdBu\", extent=extent, \n",
    "                          interpolation='nearest', vmin=-clip, vmax=clip)        \n",
    "    plt.title('Pressure wavefield')\n",
    "    plt.xlabel('x [m]')\n",
    "    plt.ylabel('z [m]')    \n",
    "    \n",
    "    # Plot Vp-model\n",
    "    ax2 = plt.subplot(122)\n",
    "    image1 = plt.imshow((vp.T)/1000, cmap=plt.cm.viridis, interpolation='nearest', \n",
    "                        extent=extent)\n",
    "    \n",
    "    plt.title('Vp-model')\n",
    "    plt.xlabel('x [m]')\n",
    "    plt.setp(ax2.get_yticklabels(), visible=False) \n",
    "    \n",
    "    divider = make_axes_locatable(ax2)\n",
    "    cax2 = divider.append_axes(\"right\", size=\"2%\", pad=0.1)\n",
    "    fig.colorbar(image1, cax=cax2)         \n",
    "    plt.ion()    \n",
    "    plt.show(block=False)\n",
    "    \n",
    "    # Calculate Partial Derivatives\n",
    "    # -----------------------------\n",
    "    for it in range(nt):\n",
    "    \n",
    "        # FD approximation of spatial derivative by 3 point operator\n",
    "        d2px, d2pz = update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz)\n",
    "\n",
    "        # Time Extrapolation\n",
    "        # ------------------\n",
    "        pnew = 2 * p - pold + vp2 * dt**2 * (d2px + d2pz)\n",
    "\n",
    "        # Add Source Term at isrc\n",
    "        # -----------------------\n",
    "        # Absolute pressure w.r.t analytical solution\n",
    "        pnew[isrc,jsrc] = pnew[isrc,jsrc] + src[it] / (dx * dz) * dt ** 2\n",
    "        \n",
    "        # Apply absorbing boundary frame\n",
    "        # ------------------------------\n",
    "        p *= absorb_coeff\n",
    "        pnew *= absorb_coeff\n",
    "        \n",
    "        # Remap Time Levels\n",
    "        # -----------------\n",
    "        pold, p = p, pnew\n",
    "    \n",
    "        # display pressure snapshots \n",
    "        if (it % isnap) == 0:\n",
    "            image.set_data(p.T)\n",
    "            fig.canvas.draw()  \n",
    "            \n",
    "    return vp        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 1: Homogeneous Model\n",
    "\n",
    "As a reference, we start with a problem, which should be quite familiar to you - the homogeneous model. I modified the FD code to define models in a separate Python function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Homogeneous model\n",
    "def model(nx,nz,vp,dx,dz):\n",
    "    \n",
    "    \n",
    "    return vp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time to define the modelling parameters and run the new FD code ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "xsrc = 1000.0 # x-source position (m)\n",
    "zsrc = xsrc   # z-source position (m)\n",
    "\n",
    "dx   = 5.0   # grid point distance in x-direction (m)\n",
    "dz   = dx     # grid point distance in z-direction (m)\n",
    "f0   = 100.0  # centre frequency of the source wavelet (Hz)\n",
    "\n",
    "# calculate dt according to the CFL-criterion\n",
    "dt = dx / (np.sqrt(2.0) * vp0)\n",
    "\n",
    "vp_hom = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 2: Random Medium\n",
    "\n",
    "Next, we add some random perturbations to the homogeneous Vp-model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Random medium model\n",
    "def model(nx,nz,vp,dx,dz):\n",
    "    \n",
    "    return vp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xsrc = 1000.0 # x-source position (m)\n",
    "zsrc = xsrc    # z-source position (m)\n",
    "\n",
    "dx   = 5.0   # grid point distance in x-direction (m)\n",
    "dz   = dx     # grid point distance in z-direction (m)\n",
    "f0   = 100.0  # centre frequency of the source wavelet (Hz)\n",
    "vpmax = vp0\n",
    "\n",
    "# calculate dt according to the CFL-criterion\n",
    "dt = dx / (np.sqrt(2.0) * vpmax)\n",
    "\n",
    "vp_rand = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 3: Fault Zone\n",
    "\n",
    "In this problem, we model acoustic wave propagation in a vertical fault zone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vertical fault zone model\n",
    "def model(nx,nz,vp,dx,dz):    \n",
    "        \n",
    "    return vp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xsrc = 1000.0 # x-source position (m)\n",
    "zsrc = 1000.0 # z-source position (m)\n",
    "\n",
    "dx   = 5.0   # grid point distance in x-direction (m)\n",
    "dz   = dx     # grid point distance in z-direction (m)\n",
    "f0   = 100.0  # centre frequency of the source wavelet (Hz)\n",
    "vpmax = 4200.0\n",
    "# vpmax = np.max(vp0)\n",
    "\n",
    "# calculate dt according to the CFL-criterion\n",
    "dt = dx / (np.sqrt(2.0) * vpmax)\n",
    "\n",
    "vp_fault = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 4: Simplified Vulcano\n",
    "\n",
    "How does the surface topography of a vulcano scatter the acoustic wavefield?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simplified vulcano (Gaussian hill)\n",
    "def model(nx,nz,vp,dx,dz):\n",
    "            \n",
    "    return vp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xsrc = 1000.0 # x-source position (m)\n",
    "zsrc = 200.0   # z-source position (m)\n",
    "\n",
    "dx   = 5.0   # grid point distance in x-direction (m)\n",
    "dz   = dx     # grid point distance in z-direction (m)\n",
    "f0   = 100.0  # centre frequency of the source wavelet (Hz)\n",
    "vpmax = np.max(vp0)\n",
    "\n",
    "# calculate dt according to the CFL-criterion\n",
    "dt = dx / (np.sqrt(2.0) * vpmax)\n",
    "\n",
    "vp_topo = FD_2D_acoustic_JIT(dt,dx,dz,f0,xsrc,zsrc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 5: Create your own problem here\n",
    "\n",
    "##### Excercise\n",
    "\n",
    "Now it's your turn, create a 2D P-wave velocity model and compute the pressure wavefield by 2D acoustic FD modelling. Pay attention to satisfy the CFL and grid dispersion criteria by choosing appropriate spatial grid point distances $dx = dz$ and time step dt."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What we learned:\n",
    "\n",
    "- How to model wave propgation in simple heterogeneous 2D acoustic media "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
